{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE EmptyDataDecls             #-}
{-# LANGUAGE TypeSynonymInstances       #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE TypeOperators              #-}
{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  Ochre.Language.Prim
-- Copyright   :  (c) Stephen Tetley 2012
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Monadic primitives, signals etc.
--
--------------------------------------------------------------------------------

module Ochre.Language.Prim
  (

    ErrMsg
  , ThrowErr(..)
  , SPair(..)
  , (:::)
  , Orch
  , execOrch
  , instr
  , globalAssignTable
  , globalAssign        -- only exported for defintion of zakinit like funs
  , globalInvoke        -- -ditto-
  , Instr
  , getInstr
  , InstDef0
  , InstDef
  , definstr
  , defvar

  , Sig
  , GenSig
  , GenSig2
  , GenSig3
  , GenSig4
  , GenISig
  , GenKSig
  , GenASig

  , (&=)
  , (&\=)
  , (&~=) 
  , defsig
  , defsig2

  , PValue(..)      -- only exported for definition of opcodeN like funs
  , Code(..)        --  -ditto-
  

  , TableIxTag
  , TableIx
 
  , ARate(..)
  , KRate(..)
  , IRate(..)

  , ASig
  , KSig
  , ISig


  , IString
  , IBool
  , IInt
  , IFloat 


  , InitValue
  , rateId
  , IKA_Signal
  , IK_Signal
  , KA_Signal


  , primInstr
  , tableIx

  , unaryOp
  , binOp

  , opcode0
  , opcode1
  , opcode2
  , opcode3
  , opcode4

  , function


  , cond
  , (.==.)
  , (./=.)
  , (.>.)
  , (.<.) 
  , (.>=.)
  , (.<=.) 

  , pfield
  , liftS

  , CastA
  , castA
  , CastK
  , castK
  , signalI
  , signalK
  , signalA

  , SMax
  , (.+.)  
  , (.-.)
  , (.*.)
  , (./.)

  , (*^)
  , (^*)
  , (/^)
  , (^/)
  , (+^)
  , (^+)
  , (-^)
  , (^-)


  ) where

import Ochre.Compiler.CompilerMon ( ErrMsg )
import Ochre.Compiler.Operators
import Ochre.Compiler.Syntax hiding ( Unit )
import Ochre.Compiler.Utils
import Ochre.Compiler.VarId


import Control.Applicative
import Data.Monoid
import Data.String



class ThrowErr m where
  throwErr :: ErrMsg -> m a


-- Strict pair with special constructor syntax

data SPair a b = !a ::: ! b 
  deriving (Eq,Ord,Show)

infixr 4 :::

type (:::) a b = SPair a b


instance (Monoid a, Monoid b) => Monoid (SPair a b) where
  mempty = mempty ::: mempty
  (a1 ::: b1) `mappend` (a2 ::: b2) = (a1 `mappend` a2) ::: (b1 `mappend` b2)




--------------------------------------------------------------------------------
-- Orch Monad allowing construction with do-notation

newtype Orch a = Orch { getOrch :: Either ErrMsg (Wo a) }

data Wo a = Wo a !(H GlobalStmt) !(H Instrument)

instance Functor Orch where
  fmap f ma = Orch $ getOrch ma >>= \(Wo a vh ih) -> return $ Wo (f a) vh ih

instance Applicative Orch where
  pure a    = Orch $ return $ Wo a emptyH emptyH
  mf <*> ma = Orch $ getOrch mf >>= \(Wo f v1 i1) ->
                     getOrch ma >>= \(Wo a v2 i2) ->
                     return $ Wo (f a) (v1 `appendH` v2) (i1 `appendH` i2)



instance Monad Orch where
  return   = pure
  ma >>= k = Orch $ getOrch ma    >>= \(Wo a v1 i1) ->
                    getOrch (k a) >>= \(Wo b v2 i2) -> 
                    return $ Wo b (v1 `appendH` v2) (i1 `appendH` i2)


instance ThrowErr Orch  where
  throwErr = Orch . Left 

execOrch :: Orch () -> Either ErrMsg Orchestra 
execOrch ma = fmap sk $ getOrch ma
  where
    sk (Wo _ vh ih) = Orchestra { orch_globals = toListH vh
                                , orch_instrs  = toListH ih
                                }

tellGlobalStmt :: GlobalStmt -> Orch () 
tellGlobalStmt stmt = Orch $ return $ Wo () (wrapH stmt) emptyH

tellInstrument :: Instrument -> Orch ()
tellInstrument inst = Orch $ return $ Wo () emptyH (wrapH inst) 


-- We use do-notation as we may have many instruments and
-- thus need need (>>)



newtype Instr = Instr { getInstr :: Int -> Either ErrMsg Instrument }

newtype InstDef0 a = InstDef0 { getInstDef0 :: (a, (H Decl)) }

type InstDef  = InstDef0 (Code ())

instance Functor InstDef0 where 
  fmap f ma = InstDef0 $ let (a,w) = getInstDef0 ma in (f a, w)

instance Applicative InstDef0 where 
  pure a = InstDef0 $ (a,emptyH)
  mf <*> ma = InstDef0 $ let (f,w1) = getInstDef0 mf
                             (a,w2) = getInstDef0 ma
                         in (f a, w1 `appendH` w2)

instance Monad InstDef0 where
  return   = pure
  ma >>= k = InstDef0 $ let (a,w1) = getInstDef0 ma
                            (b,w2) = getInstDef0 (k a)
                        in (b, w1 `appendH` w2)

instr :: Int -> Instr -> Orch ()
instr i def = either throwErr tellInstrument $ getInstr def i


primInstr :: Int -> [Decl] -> Code () -> Either ErrMsg Instrument
primInstr num vdecls ostmt = 
    fmap sk $ execCode ostmt
  where
    sk code = Instrument num vdecls (trans1 code)


globalAssignTable :: String -> RateId -> PrimId -> [Value] -> Orch TableIx
globalAssignTable ss ri pid vals = tellGlobalStmt stmt >> return (pvalue vid)
  where
    vid   = Global ri ss
    stmt  = GlobalDecl vid pid vals

globalAssign :: String -> RateId -> PrimId -> [Value] -> Orch (Sig a)
globalAssign ss ri pid vals = tellGlobalStmt stmt >> return (pvalue vid)
  where
    vid   = Global ri ss
    stmt  = GlobalDecl vid pid vals

globalInvoke :: PrimId -> [Value] -> Orch ()
globalInvoke pid vals = tellGlobalStmt stmt >> return ()
  where
    stmt = GlobalOpCall0 pid vals



defvar :: String -> PValue a -> InstDef0 (PValue a)
defvar ss val = InstDef0 $ 
    let vid = DeclVar I_RATE ss 
    in (PValue $ Var vid, wrapH $ Decl vid (getPValue val))
  


definstr :: InstDef -> Instr
definstr ma = Instr $ \i -> 
    let (code,w) = getInstDef0 ma in primInstr i (toListH w) code





--------------------------------------------------------------------------------

data OValue = Value Value
            | OpCallO PrimId [Value]
  deriving (Eq,Show)


-- Design note - it may be better not to represent a list of
-- statements with Seq and use @ type Stmts = [Stmt] @ instead.
-- This avoids nesting problems in translation - Seq allows 
-- nesting both at the left and right when we only want it at 
-- the right.

data Stmt = Seq Stmt Stmt
          | VarDef   VarId  Value                -- ihz = pfield 4
          | SigDef1  VarId  OValue               -- asig  oscil 440
          | SigDef2  VarId  VarId   OValue
          | Update   VarId  OValue
          | OpCallS  PrimId [Value]
  deriving (Eq,Show)            

trans1 :: Stmt -> Expr
trans1 (Seq (VarDef vid val)   e)   = LetValue vid val (trans1 e)
trans1 (Seq (SigDef1 vid oval) e)   = LetSig1 vid (ovalue oval) (trans1 e)
trans1 (Seq (SigDef2 v1 v2 oval) e) = LetSig2 v1 v2 (ovalue oval) (trans1 e)
trans1 (Seq (Update vid oval)  e)   = WriteSig vid (ovalue oval) :>> trans1 e
trans1 (Seq (OpCallS pid vs)   e)   = OpCall pid vs :>> trans1 e 
trans1 (OpCallS pid vs)             = OpCall pid vs
trans1 s                            = error $ "invalid syntax - " ++ show s


ovalue :: OValue -> Expr 
ovalue (Value val)      = Return val
ovalue (OpCallO pid vs) = OpCall pid vs



type Wc = H Stmt

data St = St { icount :: Int, kcount :: Int, acount :: Int }

newtype Code a = Code { getCode :: St -> Either ErrMsg (a, St, Wc) }



instance ThrowErr Code  where
  throwErr msg = Code $ \_ -> Left msg


instance Functor Code  where
  fmap f ma = Code $ \s -> getCode ma s >>= \(a,s1,w) -> return (f a,s1,w)


instance Applicative Code where
  pure a    = Code $ \s -> return (a,s,emptyH)
  mf <*> ma = Code $ \s -> getCode mf s  >>= \(f,s1,w1) ->
                           getCode ma s1 >>= \(a,s2,w2) ->
                           return (f a,s2, w1 `appendH` w2)

instance Monad Code where
  return    = pure
  ma >>= k  = Code $ \s -> getCode ma s     >>= \(a,s1,w1) -> 
                           getCode (k a) s1 >>= \(b,s2,w2) ->
                           return (b,s2, w1 `appendH` w2)


execCode :: Code a -> Either ErrMsg Stmt
execCode ma = sk =<< getCode ma (St 1 1 1)
  where
    sk (_,_,w)      = step1 $ toListH w
    step1 []        = Left "execCode - empty code"
    step1 (x:xs)    = Right $ step2 x xs
    step2 ac []     = ac
    step2 ac (x:xs) = Seq ac $ step2 x xs
   
  

newVarId :: RateId -> Code VarId
newVarId rid = Code $ \s -> return $ case rid of
    A_RATE -> (mkvar (acount s), let a = acount s in s{ acount = a+1 }, emptyH)
    K_RATE -> (mkvar (kcount s), let k = kcount s in s{ kcount = k+1 }, emptyH)
    I_RATE -> (mkvar (icount s), let i = icount s in s{ icount = i+1 }, emptyH)
  where
    mkvar i = Local rid i

tellStmt :: Stmt -> Code ()
tellStmt ast = Code $ \s -> return ((),s, wrapH ast)


-- | Phantom type wrapping Values.
--
newtype PValue a = PValue { getPValue :: Value } 

type IString        = PValue String
type IBool          = PValue Bool
type IInt           = PValue Int
type IFloat         = PValue Float


-- | Separate type for table indices (with are otherwise Int) 
-- disallowing Int operations e.g. Eq, Ord, Num.
--
-- This is analogue to disallowing pointer arithmetic.
-- 
data TableIxTag

type TableIx        = PValue TableIxTag


data IRate          = IRate deriving (Show)
data KRate          = KRate deriving (Show)
data ARate          = ARate deriving (Show)

type Sig a          = PValue a

type ISig           = Sig IRate
type KSig           = Sig KRate
type ASig           = Sig ARate

type Sig2 a         = (Sig a, Sig a)


tableIx :: Int -> TableIx
tableIx = PValue . Int



pvalue :: VarId -> PValue a 
pvalue = PValue . Signal


class InitValue ix where
   rateId :: ix -> RateId

instance InitValue ARate       where rateId _ = A_RATE
instance InitValue KRate       where rateId _ = K_RATE
instance InitValue IRate       where rateId _ = I_RATE
instance InitValue TableIx     where rateId _ = I_RATE


class InitValue rate => IKA_Signal rate 

instance IKA_Signal ARate
instance IKA_Signal KRate
instance IKA_Signal IRate

class IK_Signal a

instance IK_Signal KRate
instance IK_Signal IRate

class KA_Signal a

instance KA_Signal KRate
instance KA_Signal ARate




instance IsString IString where
  fromString = PValue . String


instance Num IInt where
  (+) = binOp ADD
  (-) = binOp SUB
  (*) = binOp MUL
  abs _     = error $ "abs - no implementation for VInt"
  signum _  = error $ "signum - no implementation for VInt"
  fromInteger = PValue . Int . fromInteger


instance Num IFloat where
  (+) = binOp ADD
  (-) = binOp SUB
  (*) = binOp MUL
  abs _     = error $ "abs - no implementation for VFloat"
  signum _  = error $ "signum - no implementation for VFloat"
  fromInteger = PValue . Int . fromInteger


instance Num ISig where
  (+) = binOp ADD
  (-) = binOp SUB
  (*) = binOp MUL
  abs _     = error $ "abs - no implementation for ISig"
  signum _  = error $ "signum - no implementation for ISig"
  fromInteger = PValue . Float . fromInteger


instance Num KSig where
  (+) = binOp ADD
  (-) = binOp SUB
  (*) = binOp MUL
  abs _     = error $ "abs - no implementation for KSig"
  signum _  = error $ "signum - no implementation for KSig"
  fromInteger = PValue . Float . fromInteger


instance Num ASig where
  (+) = binOp ADD
  (-) = binOp SUB
  (*) = binOp MUL
  abs _     = error $ "abs - no implementation for ASig"
  signum _  = error $ "signum - no implementation for ASig"
  fromInteger = PValue . Float . fromInteger


instance Fractional IFloat where
  (/)       = binOp DIV
  recip _   = error $ "recip - no implementation for VFloat"
  fromRational = PValue . Float . fromRational


instance Fractional ISig where
  (/)       = binOp DIV
  recip _   = error $ "recip - no implementation for ISig"
  fromRational = PValue . Float . fromRational


instance Fractional KSig where
  (/)       = binOp DIV
  recip _   = error $ "recip - no implementation for KSig"
  fromRational = PValue . Float . fromRational

instance Fractional ASig where
  (/)       = binOp DIV
  recip _   = error $ "recip - no implementation for ASig"
  fromRational = PValue . Float . fromRational




-- Can define NUM instances for empty types as PExpr is a 
-- Phantom.

unaryOp :: UnaryOp -> PValue a -> PValue a
unaryOp op a = PValue $ 
    UnaryV op (getPValue a)




binOp :: BinOp -> PValue a -> PValue a -> PValue a
binOp op a b = PValue $ 
    BinV op (getPValue a) (getPValue b)


relOp :: RelOp -> PValue a -> PValue a -> PValue Bool
relOp op a b = PValue $ 
    RelV op (getPValue a) (getPValue b)



-- | The return type of a opcode is not actually a signal, it
-- is /code that generates a signal/.
--
newtype GenSig rate = GenSig { getGenSig :: Code OValue }


type GenISig = GenSig IRate
type GenKSig = GenSig KRate
type GenASig = GenSig ARate

type GenSig2 a = GenSig (a,a)
type GenSig3 a = GenSig (a,a,a)
type GenSig4 a = GenSig (a,a,a,a)



defsig :: IKA_Signal rate => rate -> GenSig rate -> Code (Sig rate)
defsig rid ma = do
    vid     <- newVarId $ rateId rid
    code    <- getGenSig ma
    tellStmt $ SigDef1 vid code
    return $ PValue $ Signal vid


defsig2 :: IKA_Signal rate => rate -> GenSig2 rate -> Code (Sig2 rate)
defsig2 rid ma = do
    v1      <- newVarId $ rateId rid
    v2      <- newVarId $ rateId rid
    code    <- getGenSig ma
    tellStmt $ SigDef2 v1 v2 code
    return (PValue $ Signal v1, PValue $ Signal v2)


infixr 1 &=

-- | Assignment
--
-- Note - eliminates redundant assignment, i.e
--
-- > asig1 &= asig1 
--
(&=) :: Sig rate -> GenSig rate -> Code ()
(&=) (PValue val) ma = case val of
    Signal vid -> do { code <- getGenSig ma 
                     ; if code /= (Value val) then tellStmt $ Update vid code
                                              else return () }
    _          -> throwErr "Assignment - lhs not a signal (lvar)."


--
-- Note - putting the slash as the prefix makes user code look
-- horrible when lambdas are used. 
--

infixr 1 &\=

(&\=) :: Sig rate -> (Sig rate -> Sig rate) -> Code ()
(&\=) sig fn = sig &= liftS (fn sig)



infixr 1 &~=
   
(&~=) :: Sig rate -> Sig rate -> Code ()
(&~=) sig rhs = sig &= liftS rhs
   
 
   
-- Opcodes must return code, rather than /tell/ code. 
-- The exception seems to be opcode0 e.g. out, outs
--
opcode0 :: String -> [Value] -> Code ()
opcode0 ss vs = tellStmt $ OpCallS (PrimId ss) vs


opcode1 :: String -> [Value] -> GenSig rate
opcode1 ss vs = GenSig $ return $ OpCallO (PrimId ss) vs

opcode2 :: String -> [Value] -> GenSig2 rate
opcode2 ss vs = GenSig $ return $ OpCallO (PrimId ss) vs

opcode3 :: String -> [Value] -> GenSig3 rate
opcode3 ss vs = GenSig $ return $ OpCallO (PrimId ss) vs

opcode4 :: String -> [Value] -> GenSig4 rate
opcode4 ss vs = GenSig $ return $ OpCallO (PrimId ss) vs




type family SMax a b :: *

type instance SMax ARate ARate = ARate
type instance SMax ARate KRate = ARate
type instance SMax KRate ARate = ARate
type instance SMax ARate IRate = ARate
type instance SMax IRate ARate = ARate

type instance SMax KRate KRate = KRate
type instance SMax KRate IRate = KRate
type instance SMax IRate KRate = KRate

type instance SMax IRate IRate = IRate


infixl 6 .+.

(.+.) :: c ~ SMax a b => Sig a -> Sig b -> Sig c
(.+.) s1 s2 = PValue $ BinV ADD (getPValue s1) (getPValue s2)


infixl 6 .-.

(.-.) :: c ~ SMax a b => Sig a -> Sig b -> Sig c
(.-.) s1 s2 = PValue $ BinV SUB (getPValue s1) (getPValue s2)

infixl 7 .*.

(.*.) :: c ~ SMax a b => Sig a -> Sig b -> Sig c
(.*.) s1 s2 = PValue $ BinV MUL (getPValue s1) (getPValue s2)

infixl 7 ./.

(./.) :: c ~ SMax a b => Sig a -> Sig b -> Sig c
(./.) s1 s2 = PValue $ BinV DIV (getPValue s1) (getPValue s2)



cond :: IBool -> PValue a -> PValue a -> PValue a
cond test a b = PValue $ CondV (getPValue test) (getPValue a) (getPValue b)

infix 4 .==.

(.==.) :: Eq a => PValue a -> PValue a -> IBool
(.==.) = relOp EQU


infix 4 ./=. 

(./=.) :: Eq a => PValue a -> PValue a -> IBool
(./=.) = relOp NEQU


infix 4 .>. 

(.>.) :: Eq a => PValue a -> PValue a -> IBool
(.>.) = relOp GREATTHAN

infix 4 .<. 

(.<.) :: Eq a => PValue a -> PValue a -> IBool
(.<.) = relOp LESSTHAN

infix 4 .>=. 

(.>=.) :: Eq a => PValue a -> PValue a -> IBool
(.>=.) = relOp GTEQU

infix 4 .<=. 

(.<=.) :: Eq a => PValue a -> PValue a -> IBool
(.<=.) = relOp LTEQU

pfield :: Int -> PValue a
pfield = PValue . Pfield

liftS :: Sig a -> GenSig a
liftS s1 = GenSig $ return $ Value $ getPValue s1

function :: String -> [Value] -> PValue ans
function name vs = PValue $ FunCallV (PrimId name) vs



signalI :: IFloat -> Sig IRate
signalI = PValue . getPValue

signalK :: IFloat -> Sig KRate
signalK = castK . signalI

signalA :: IFloat -> Sig ARate
signalA = castA . signalI


class CastK a where castK :: Sig a -> Sig KRate

instance CastK IRate where
  castK = PValue . getPValue

instance CastK KRate where
  castK = id


class CastA a where castA :: Sig a -> Sig ARate

instance CastA IRate where
  castA = PValue . getPValue

instance CastA KRate where
  castA = PValue . getPValue

instance CastA ARate where
  castA = id



infixr 7 *^
infixl 7 ^*

-- | Multiply a signal by an I-rate constant.
--
(*^) :: (IKA_Signal a, Num (Sig a)) => IFloat -> Sig a -> Sig a
(*^) d s = (cast d :: Sig a) * s
  where 
    cast = PValue . getPValue

(^*) :: (IKA_Signal a, Num (Sig a)) => Sig a -> IFloat -> Sig a
(^*) s d = s * (cast d :: Sig a) 
  where 
    cast = PValue . getPValue



infixr 7 /^
infixl 7 ^/
-- | Divide a signal by an I-rate constant.
--
(/^) :: (IKA_Signal a, Fractional (Sig a)) => IFloat -> Sig a -> Sig a
(/^) d s = (cast d :: Sig a) / s
  where 
    cast = PValue . getPValue

(^/) :: (IKA_Signal a, Fractional (Sig a)) => Sig a -> IFloat -> Sig a
(^/) s d = s / (cast d :: Sig a)
  where
    cast = PValue . getPValue


infixr 6 +^
infixl 6 ^+

-- | Addition of a signal and an I-rate constant.
--
(+^) :: (IKA_Signal a, Num (Sig a)) => IFloat -> Sig a -> Sig a
(+^) d s = (cast d :: Sig a) + s
  where 
    cast = PValue . getPValue

(^+) :: (IKA_Signal a, Num (Sig a)) => Sig a -> IFloat -> Sig a
(^+) s d = s + (cast d :: Sig a) 
  where
    cast = PValue . getPValue


infixr 6 -^
infixl 6 ^-

-- | Subtraction from a signal by an I-rate constant.
--
(-^) :: (IKA_Signal a, Num (Sig a)) => IFloat -> Sig a -> Sig a
(-^) d s = (cast d :: Sig a) - s
  where 
    cast = PValue . getPValue



(^-) :: (IKA_Signal a, Num (Sig a)) => Sig a -> IFloat -> Sig a
(^-) s d = s - (cast d :: Sig a)
  where 
    cast = PValue . getPValue


